package ltd.hello.trik.snake;

import controlP5.Button;
import controlP5.ControlP5;
import peasy.PeasyCam;
import processing.core.PApplet;
import processing.data.Table;

import java.io.File;
import java.util.ArrayList;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

import static ltd.hello.trik.snake.SnakeCommon.*;

/**
 * +----------------------------------------------------------------------
 * |
 * +----------------------------------------------------------------------
 * | <a href="https://www.hellotrik.ltd">Copyright (c) 2024~2025 </a>.
 * +----------------------------------------------------------------------
 * | Licensed MIT
 * +----------------------------------------------------------------------
 * | 作者: 🦠 matrixg 👾 <hellotrik@foxmail.com>
 * | 日期: 2024/9/6
 * +----------------------------------------------------------------------
 */
public class SnakeAI extends PApplet {
    static ControlP5 cp5;
    Button graphButton, loadButton, saveButton, increaseMut, decreaseMut;
    Population pop;

    public void Graph(int theValue) {
        println("a button event from colorA: " + theValue);
        SnakeCommon.graph = new EvolutionGraph();
    }

    public void Load() {
        selectInput("Load Snake Model", "fileLoad");
    }

    public void Save() {
        selectOutput("Save Snake Model", "fileSelectedOut");
    }

    public void add() {
        SnakeCommon.mutationRate *= 2;
        SnakeCommon.defaultmutation = SnakeCommon.mutationRate;
    }

    public void sub() {
        SnakeCommon.mutationRate /= 2;
        SnakeCommon.defaultmutation = SnakeCommon.mutationRate;
    }

    public void settings() {
        size(1920, 1280, P3D);
    }

    public void setup() {
        cam = new PeasyCam(this, 1920);
        SnakeCommon.font = createFont("agencyfb-bold.ttf", 32);
        SnakeCommon.evolution = new ArrayList<>();
        cp5 = new ControlP5(this);
        cp5.setAutoDraw(false);
        loadButton = cp5.addButton("Load").setVisible(false).setPosition(120, 5).setSize(80, 20);
        saveButton = cp5.addButton("Save").setVisible(false).setPosition(210, 5).setSize(80, 20);
        graphButton = cp5.addButton("Graph").setVisible(false).setPosition(300, 5).setSize(80, 20);
        increaseMut = cp5.addButton("add").setVisible(false).setLabel("+").setPosition(330, 95).setSize(15, 15);
        decreaseMut = cp5.addButton("sub").setVisible(false).setLabel("-").setPosition(355, 95).setSize(15, 15);
        graphButton.show();
        loadButton.show();
        saveButton.show();
        frameRate(SnakeCommon.fps);
        pop = new Population(1000, this); //adjust size of population
//        Load();
        modelLoaded=true;
    }

    public void draw() {
        background(0);
        noFill();
        stroke(255);
        rectMode(CORNER);
        textFont(SnakeCommon.font);
        box(height-5*SIZE);
        if(modelLoaded) {
            if (pop.done()) {
                SnakeCommon.highscore = pop.bestSnake.score;
                pop.calculateFitness();
                pop.naturalSelection();
            } else {
                pop.update();
                pop.show();
            }
        }
        cam.beginHUD();
        textSize(16);
        textAlign(LEFT);
        fill(150);
        text("BEST FITNESS : " + pop.bestFitness, 120, 50);
        text("GEN : " + pop.gen, 120, 70);
        text("MOVES LEFT : " + pop.bestSnake.lifeLeft, 120, 90);
        text("MUTATION RATE : " + SnakeCommon.mutationRate * 100 + "%", 120, 110);
        text("SCORE : " + pop.bestSnake.score, 120, height - 65);
        text("LIVE : " + pop.getLives(), 120, height - 45);
        text("HIGHSCORE : " + SnakeCommon.highscore, 120, height - 25);
        increaseMut.show();
        decreaseMut.show();
        fill(255, 0, 0);
        text("RED < 0", 120, height-5);
        fill(0, 0, 255);
        text("BLUE > 0", 180, height-5);
        cp5.draw(); // 绘制 ControlP5 的 UI 控件
        cam.endHUD();
    }

    public boolean overControlP5() {
        return cp5.isMouseOver(); // 检查鼠标是否在 ControlP5 的控件上
    }

    public void fileLoad(File selection) {
        if (selection == null) {
            println("Window was closed or the user hit cancel.");
        } else {
            var path = selection.getAbsolutePath();
            var modelTable = loadTable(path, "header");
            var weights = new Matrix[modelTable.getColumnCount() - 1];
            var in = new float[hidden_nodes][25];
            IntStream.range(0, hidden_nodes).forEach(i -> IntStream.range(0,25).forEach(j->in[i][j] = modelTable.getFloat(j + i * 25, "L0")));
            weights[0] = new Matrix(in);
            IntStream.range(1, weights.length-1).forEach(h->{
                float[][] hid = new float[hidden_nodes][hidden_nodes + 1];
                IntStream.range(0,hidden_nodes).forEach(i->IntStream.range(0,hidden_nodes+1).forEach(j->hid[i][j] = modelTable.getFloat(j + i * (hidden_nodes + 1), "L" + h)));
                weights[h] = new Matrix(hid);
            });
            var out = new float[4][hidden_nodes + 1];
            IntStream.range(0,4).forEach(i->IntStream.range(0,hidden_nodes+1).forEach(j->out[i][j] = modelTable.getFloat(j + i * (hidden_nodes + 1), "L" + (weights.length - 1))));
            weights[weights.length - 1] = new Matrix(out);
            pop.load(weights);
            modelLoaded = true;
            evolution = new ArrayList<>();
            int g = 0;
            int genscore = modelTable.getInt(g, "Graph");
            while (genscore != 0) {
                evolution.add(genscore);
                g++;
                genscore = modelTable.getInt(g, "Graph");
            }
        }
    }

    public void fileSelectedOut(File selection) {
        if (selection == null) {
            println("Window was closed or the user hit cancel.");
        } else {
            var path = selection.getAbsolutePath();
            var modelTable = new Table();
            var modelToSave = pop.bestSnake.clone();
            var modelWeights = modelToSave.brain.pull();
            var weights = new float[modelWeights.length][];
            IntStream.range(0, weights.length).parallel().forEach(i -> {
                weights[i] = modelWeights[i].toArray();
            });
            IntStream.range(0, weights.length).forEach(i -> {
                modelTable.addColumn("L" + i);
            });
            modelTable.addColumn("Graph");
            int maxLen = IntStream.range(0, weights.length)
                    .parallel()
                    .map(i -> weights[i].length)  // 映射到每个权重的长度
                    .reduce(0, Math::max);

            var g = new AtomicInteger(0);
            IntStream.range(0,maxLen).forEach(i->{
                var newRow = modelTable.addRow();
                IntStream.range(0,weights.length+1).parallel().forEach(j -> {
                    if (j == weights.length) {
                        int gVal = g.getAndIncrement();
                        if (gVal < evolution.size()) {
                            newRow.setInt("Graph", evolution.get(gVal));
                        }
                    } else if (i < weights[j].length) {
                        newRow.setFloat("L" + j, weights[j][i]);
                    }
                });
            });
            saveTable(modelTable, path);
        }
    }

    public void mousePressed() {
        // 否则启用 PeasyCam
        cam.setActive(!overControlP5()); // 当鼠标在 UI 控件上时禁用 PeasyCam
    }

}
